import json
import base64
import requests
from jose import jwt


class OIDC:
    def __init__(self, **kwargs):
        self.client_id = ''
        self.client_secret = ''
        self.access_url = ''
        self.verify_url = ''
        self.redirect_url = ''
        self.cert = ''
        self.uid = 'preferred_username'
        self.name = 'name'
        if kwargs:
            self.init(**kwargs)

    def init(self, client_id=None, client_secret=None, access_url=None, verify_url=None, redirect_url=None, cert='',
             uid='preferred_username', name: str = 'name'):
        self.client_id = client_id
        self.client_secret = client_secret
        self.access_url = access_url
        self.verify_url = verify_url
        self.redirect_url = redirect_url
        self.cert = cert
        self.uid = uid
        self.name = name

    def _verify_token(self, token: dict):
        """token 验证"""

        if not self.verify_url:
            return

        auth = f'{self.client_id}:{self.client_secret}'.encode('UTF-8')
        header = {
            'Content-Type': 'application/x-www-form-urlencoded',
            'Authorization': 'Basic ' + base64.b64encode(auth).decode('UTF-8')
        }
        data = {
            'token_type_hint': 'access_token',
            'token': token['access_token']
        }
        resp = requests.post(url=self.verify_url, data=data, headers=header)
        if resp.status_code != 200:
            raise Exception('验证失败')

        result: dict = json.loads(resp.text)
        is_active = result['active'] if 'active' in result else True
        if not is_active:
            raise Exception(f'账号已禁用')

    def get_token(self, code: str) -> str:
        """获取token"""
        data = {
            'grant_type': 'authorization_code',
            'client_id': self.client_id,
            'client_secret': self.client_secret,
            'code': code,
            'redirect_uri': self.redirect_url,  # 针对 keycloak, 正常不会使用
        }

        resp = requests.post(url=self.access_url, data=data,
                             headers={'Content-Type': 'application/x-www-form-urlencoded'})
        if resp.status_code != 200:
            raise Exception('获取Token失败')

        token = json.loads(resp.text)
        self._verify_token(token)
        return token['access_token']

    def get_user_info_by_token(self, access_token: str) -> dict:
        try:
            option = {
                'verify_signature': True if self.cert else False,
                'verify_aud': False
            }

            info = jwt.decode(access_token, self.cert, audience=self.client_id, options=option)
        except Exception as e:
            raise Exception('解析Token失败')
        else:
            if self.uid not in info:
                raise Exception(f'唯一标识字段{self.uid}不存在')
            if self.name not in info:
                raise Exception(f'用户名显示字段{self.name}不存在')

            return {
                'key': info[self.uid],
                'name': info[self.name]
            }


if __name__ == '__main__':
    oidc = OIDC()

    code = ''
    token = oidc.get_token(code)
    user = oidc.get_user_info_by_token(token)
    print(user)
